#include "StdAfx.h"

#include "SecrecyAES.h"

#include "AESDefs.h"
#include "AESUtils.h"

#include "AESSubstitution.h"
#include "AESPermutation.h"

#include <malloc.h>

/**
 *  Module Name: 	SecrecyAES.cpp
 *  
 *  Module represents the core engine of the AES secrecy algorithm.
 *
 *  Defines the tranforms applied to the data and the number of 
 *  rounds. 
 * 
 * @author Ravi Mohandas, Mar 06 2011
 */

	/**
	 *	Module scoped private apis
	 * 
	 */

	static void		setupState(KeySize	keySize);
	static bool		validateState(void);
	static byte*	newKey (void);
	static int		deleteKey (void);
	static int		copyKey (const byte keyBytes[]);
	static void		transform_SubBytes (byte** stateTable, int nRows, int nCols);
	static byte*	newKeyset(void);
	static byte*	deleteKeyset(void);
	static int		expandKey (void);
	static void		applyCorekeySchedule(void);

	/**
	 *	Module scoped variables 
	 * 
	 */

	static AESConfiguration	_configuration;					// aes system params
	static byte*	_keyBytes;								// seed key bytes
	static byte		_stateTable[_NROWS][_NCOLS];			// state table for AES
	static byte		_expandedKey[_NROWS_EXP][_NCOLS_EXP];	// array holds the expanded keyset for all rounds

	extern "C"{
	/**
	 *	Configures the strength of the secrecy operations 
	 *	performed on a message. 
	 *  
	 * 
	 */
SECRECY_SYSTEM_API	
int InitAES (const byte keyBytes [], KeySize keysize)
	{	
		int Results = E_SUCCESS;

		setupState (keysize);

		// makes a copy of the seed key provided by user
		Results = 
			copyKey (keyBytes);

		// expands the key into a set of keys to use for 
		// the n rounds of tranforms to be applied
		Results = 
			expandKey ();

		return E_SUCCESS;
	}


SECRECY_SYSTEM_API	
int ShutdownAES (void)
	{	
		deleteKey ();	

		return E_SUCCESS;
	}



	/**
	 * Enciphers the user message and returns a cipher
	 * 
	 * Allocate and pass valid buffers for both input and output.
	 *
	 */
SECRECY_SYSTEM_API	
int Encrypt (const byte message[],  byte cipher[], int length){
		
		if ( !validateState () )
		{
			perror (E_INVALID_STATE);
		}


		return E_SUCCESS;
	}

	/**
	 * deciphers the cipher and returns the original message
	 * 
	 * Allocate and pass valid buffers for both input and output.
	 * 
	 */
SECRECY_SYSTEM_API	
int Decrypt (const byte cipher[], byte message[], int length){

		if ( !validateState () )
		{
			perror (E_INVALID_STATE);
		}

	
	return E_SUCCESS;
	}

	}

	/**
	 * Private Functions 
	 * 
	 * 
	 *
	 */
	
	static void setupState(KeySize	keySize)
	{
	  set_nDataChunks	(&_configuration, _FOUR_WORDS);


	  if (keySize == System_key128)
	  {
		  set_nKeyChunks (&_configuration,_FOUR_WORDS);
		  set_nRounds (&_configuration,_TEN_ROUNDS);
	  }
	  else if (keySize == System_key192)
	  {
		  set_nKeyChunks (&_configuration,_SIX_WORDS);
		  set_nRounds (&_configuration,_TWELVE_ROUNDS);
	  }
	  else if (keySize == System_key256)
	  {
		  set_nKeyChunks (&_configuration,_EIGHT_WORDS);
		  set_nRounds (&_configuration,_FOURTEEN_ROUNDS);
	  }
	  else{
		  // defaults 
		  set_nKeyChunks (&_configuration,_FOUR_WORDS);
		  set_nRounds (&_configuration,_TEN_ROUNDS);
	  }
	}

	static bool validateState(void)
	{
		return true;		
	}


	/**
	 * Copy the Key
	 * 
	 */


static int copyKey (const byte keyBytes[])
{
	int Results		= E_SUCCESS;
	int nKeyChunks	= 0;
	size_t size = 0;

	if (NULL == keyBytes ) {
		return E_INVALID_PARAMS;
	}

	_keyBytes = newKey();

	if (_keyBytes )
	{
		nKeyChunks = get_nKeyChunks (&_configuration);
		size = (_FOUR_WORDS * nKeyChunks);
		memcpy (_keyBytes, keyBytes,  size);
	}
	else{
		Results = E_NO_MEMORY;
	}

	return Results;
}


static byte* newKey (void)
{
	int Results = E_SUCCESS;
	int nKeyChunks = 0;

	nKeyChunks = get_nKeyChunks (&_configuration);
	_keyBytes = (byte*) malloc (_FOUR_WORDS * nKeyChunks);

	if (NULL == _keyBytes){
		perror (Results = E_NO_MEMORY);
	}
	else{
		memset (_keyBytes, '\0', (_FOUR_WORDS * nKeyChunks));
	}

	return _keyBytes;
}


static int deleteKey (void)
{
	int Results = E_SUCCESS;

	if (_keyBytes){

		free ( _keyBytes);

		_keyBytes = NULL;
	}

	return Results;
}


	/**
	 * 
	 * Subs for each of the values in the state table
	 *  
	 * The tranform provides the confusion component during encryption. 
	 * It increases the distance between the key used and the cipher text 
	 * produced. 
	 * 
	 * 
	 */

void transform_SubBytes (byte** stateTable, int nRows, int nCols)
{
	int iIndex = 0;
	int jIndex = 0;

	if (NULL == stateTable)
	{
		perror (E_INVALID_PARAMS);
	}

	for (iIndex = 0; iIndex < nRows; iIndex++){

		subBytes ( &_stateTable[iIndex][0], nCols);
	}
}


	/**
	 * 
	 * byte array allocated to hold the expanded key  
	 *  
	 * 
	 */

	/**
	 * 
	 * Copies the first 32 bytes of cipher key to expanded key.
	 *  
	 * 
	 */


static void copySeedKey()
{
	int row = 0;

	int nKeyChunks = get_nKeyChunks (&_configuration);

	for (row = 0; row < nKeyChunks; ++row)
	  {
		_expandedKey[row][0] = _keyBytes[4*row];
		_expandedKey[row][1] = _keyBytes[4*row+1];
		_expandedKey[row][2] = _keyBytes[4*row+2];
		_expandedKey[row][3] = _keyBytes[4*row+3];
	  }
}

	/**
	 * 
	 * Core routine to expand the seeded key.
	 *  
	 * The transform permutes and substitutes to generate an 
	 * expanded keyset. The operations used in this transform 
	 * otherthan XOR are,  
	 * 1. RotWord that is a simple rotation to the left by one byte
	 * 2. SubWord, that subs from the sBox
	 * 3. Rcon, that converts a byte doing a special arithmentic in a 
	 *    restricted field of numbers.
	 *
	 */

static void applyCorekeySchedule()
{
  int row = 0;
  int Nk = get_nKeyChunks (&_configuration);
  int Nb = get_nDataChunks(&_configuration);
  int Nr = get_nRounds(&_configuration);

  byte temp [_FOUR_BYTES];
  
  for (int row = Nk; row < Nb * (Nr+1); ++row)
  {
    temp[0] = _expandedKey[row-1][0]; 
	temp[1] = _expandedKey[row-1][1];
    temp[2] = _expandedKey[row-1][2]; 
	temp[3] = _expandedKey[row-1][3];
  
    if (row % Nk == 0)  
    {
	  RotWord(temp);
      subWord(temp);
  
      temp[0] = (byte)( (int)temp[0] ^ (int)RjConstant(row/Nk,0 ));
      temp[1] = (byte)( (int)temp[1] ^ (int)RjConstant(row/Nk,1 ));
      temp[2] = (byte)( (int)temp[2] ^ (int)RjConstant(row/Nk,2 ));
      temp[3] = (byte)( (int)temp[3] ^ (int)RjConstant(row/Nk,3 ));
    }
    else if ( Nk > 6 && (row % Nk == 4) )  
    {
	 // adds an extra sbox transformation to each new byte, in case the key is 256 bits long key
      subWord(temp);
    }
  
	// XOR temp with the 4-byte block 32 bytes before the end of the current expanded key.
    // These 4 bytes become the next bytes in the expanded key

    _expandedKey[row][0] = (byte) ( (int)_expandedKey[row-Nk,0] ^ (int)temp[0] );
    _expandedKey[row][1] = (byte) ( (int)_expandedKey[row-Nk,1] ^ (int)temp[1] );
    _expandedKey[row][2] = (byte) ( (int)_expandedKey[row-Nk,2] ^ (int)temp[2] );
    _expandedKey[row][3] = (byte) ( (int)_expandedKey[row-Nk,3] ^ (int)temp[3] );
  
  }  // for loop
}

	/**
	 * 
	 * Produces the required number of round keys using the initial seed key as input. 
	 *  
	 * The tranform selects keys from the key space keeping the probability of
	 * selection for each key selected about the same. 
	 * 
	 * The number of rounds to apply the AES transforms is a function of the selected 
	 * strength of the AES algorithm. 
	 * For e.g., AES specifies three key lengths, 128bits, 192bits and 256bits.
	 * The number of rounds of transformations applied for a 128bit key is 10 rounds. 
	 * It goes up by 2 more rounds each for 192bit and 256bit keys.
	 * 
	 */

int expandKey (void)
{
	int Results = E_SUCCESS;

	// Expanded key copies the first 32-bytes of the seed key 
	copySeedKey();

	// Walk the expanded array and for every 32 bytes (or a row of four bytes), apply core schedule
	applyCorekeySchedule();

	return Results;
}
